import os
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms


def diseases_recognition(img_path):
    idx_to_labels = np.load('idx_to_labels.npy', allow_pickle=True).item()
    # 测试集图像预处理-RCTN：缩放、裁剪、转 Tensor、归一化
    test_transform = transforms.Compose([transforms.Resize(256),
                                         transforms.CenterCrop(224),
                                         transforms.ToTensor(),
                                         transforms.Normalize(
                                             mean=[0.485, 0.456, 0.406],
                                             std=[0.229, 0.224, 0.225])
                                         ])
    print(torch.__version__)
    print(torch.cuda.is_available())
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('device:', device)
    # 定义文件路径
    file_path = "checkpoints/pest30_pytorch_20230130.pth"
    model = torch.load(file_path, map_location=device, weights_only=False)
    model = model.eval().to(device)
    img_pil = Image.open(img_path).convert('RGB')
    input_img = test_transform(img_pil)  # 预处理
    input_img = input_img.unsqueeze(0).to(device)
    # 执行前向预测，得到所有类别的 logit 预测分数
    pred_logits = model(input_img)
    pred_softmax = F.softmax(pred_logits, dim=1)  # 对 logit 分数做 softmax 运算

    n = 5
    top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
    pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
    confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度
    rec_result=[]
    for i in range(n):
        class_name = idx_to_labels[pred_ids[i]] # 获取类别名称
        confidence = "{:.2f}".format(confs[i]) # 获取置信度
        tmp={"NameZHCN":class_name,"Synonyms":confidence}
        rec_result.append(tmp)
    return rec_result